﻿#include "decoder.h"
#include <cstdio>
#include <chrono>
#include <iostream>
#include <boost/scope_exit.hpp>
#include <netinet/in.h>


const uint8_t H264_START_SEQUENCE[] = { 0, 0, 0, 1 };

Decoder::Decoder()
{
}

Decoder::~Decoder()
{
    if (m_av_packet) {
      av_packet_free(&m_av_packet);
    }
    avformat_close_input(&m_format_context);
}

void Decoder::set_mp4_path(const std::string &path)
{
    m_mp4_path = path;
}

bool Decoder::start()
{
    std::cout << "Setting up decoder" << std::endl;
    m_packet_buffer = std::make_unique<char[]>(m_packet_buffer_size);
    int ret = mpp_packet_init(&m_packet, m_packet_buffer.get(), m_packet_buffer_size);
    if (ret) {
        std::cout << "Failed to initialise MPP packet" << std::endl;
        return false;
    }

    if (!init_mpp())
        return false;

    // Open the file
    std::cout << "Opening the file: " << m_mp4_path << std::endl;
    if (avformat_open_input(&m_format_context, m_mp4_path.c_str(), nullptr, nullptr) != 0)
    {
      std::cout << "Failed to open avformat file: " << m_mp4_path << std::endl;
      return false;
    }
    if (avformat_find_stream_info(m_format_context, nullptr) < 0)
    {
      std::cout << "Failed to find stream info: " << m_mp4_path << std::endl;
      return false;
    }
    for (unsigned int i = 0; i < m_format_context->nb_streams; i++)
    {
      if ((m_format_context->streams[i]->codecpar->codec_type == AVMEDIA_TYPE_VIDEO) && (m_format_context->streams[i]->codecpar->codec_id == AVCodecID::AV_CODEC_ID_H264)) // Currently only support H264
      {
        m_videostream = i;
        break;
      }
    }
    if (!m_videostream.has_value())
    {
      std::cout << "Failed to find video stream: " << m_mp4_path << std::endl;
      return false;
    }

    // Find SPS/PPS if available and pass it to the decoder
    std::vector<uint8_t> spspps;
    if (m_format_context->streams[*m_videostream]->codecpar->extradata && m_format_context->streams[*m_videostream]->codecpar->extradata_size)
    {
      const std::vector<uint8_t> extradata(m_format_context->streams[*m_videostream]->codecpar->extradata, m_format_context->streams[*m_videostream]->codecpar->extradata + m_format_context->streams[*m_videostream]->codecpar->extradata_size);
      if (extradata.size())
      {
        if (extradata[0] >= 1) // SPS+PPS count, but we only care about the first one
        {
          const int spscount = extradata[5] & 0x1f;
          const int spsnalsize = (extradata[6] << 8) | extradata[7];
          if ((spsnalsize + 8) <= extradata.size())
          {
            std::cout << "Gathering SPS: " << spsnalsize << std::endl;
            spspps.insert(spspps.end(), extradata.data() + 8, extradata.data() + 8 + spsnalsize);
            if ((spsnalsize + 8 + 1) <= extradata.size())
            {
              const int ppscount = extradata[8 + spsnalsize] & 0x1f;
              if (ppscount >= 1)
              {
                if ((spsnalsize + 8 + 1 + 2) < extradata.size())
                {
                  const int ppsnalsize = (extradata[8 + spsnalsize + 1] << 8) | extradata[8 + spsnalsize + 2];
                  if ((spsnalsize + 8 + 1 + 2 + ppsnalsize) <= extradata.size())
                  {
                    std::cout << "Gathering PPS: " << ppsnalsize << std::endl;
                    spspps.insert(spspps.end(), H264_START_SEQUENCE, H264_START_SEQUENCE + sizeof(H264_START_SEQUENCE));
                    spspps.insert(spspps.end(), extradata.data() + 8 + spsnalsize + 3, extradata.data() + 8 + spsnalsize + 3 + ppsnalsize);
                  }
                }
              }
            }
          }
        }
      }
    }
    if (spspps.size())
    {
      std::cout << "Sending SPS and PPS" << std::endl;
      if (!SendFrame(m_mpi, m_ctx, spspps.data(), spspps.size(), m_packet))
      {
        std::cout << "Failed to send SPS+PPS frame" << std::endl;
        exit(1);
      }
    }

#if 0
    m_input_thread.reset(new std::thread([&]() {
        read_file_loop_and_do();
    }));

    m_output_thread.reset(new std::thread([&]() {
        output_loop();
    }));
#endif
    return true;
}

bool Decoder::send_frame()
{
    std::chrono::steady_clock::time_point start = std::chrono::steady_clock::now();
    static double time_base = static_cast<double>(m_format_context->streams[*m_videostream]->time_base.num) / static_cast<double>(m_format_context->streams[*m_videostream]->time_base.den);
    // Calculate time of frame
    if (m_av_packet) {
      const double m_av_packet_time = static_cast<double>(m_av_packet->pts) * time_base;
      const std::chrono::steady_clock::time_point now = std::chrono::steady_clock::now();
      if (now >= (start + std::chrono::milliseconds(static_cast<int>(m_av_packet_time * 1000.0)))) {
        av_packet_free(&m_av_packet);
//        std::cout << "av_packet free 1" << std::endl;
        m_av_packet = nullptr;
      }
    }
    // Read frame from file
    if (m_av_packet == nullptr) {
      m_av_packet = av_packet_alloc();
      int ret = av_read_frame(m_format_context, m_av_packet);
      if (ret == AVERROR_EOF) {
        ret = av_seek_frame(m_format_context, *m_videostream, 0, AVSEEK_FLAG_ANY);
        if (ret < 0) {
          std::cout << "Failed to seek frame" << std::endl;
          exit(1);
        }
        start = std::chrono::steady_clock::now();
      } else if (ret) {
        std::cout << "Failed to seek frame" << std::endl;
        exit(1);
      }
      static int i = 4;
      if (i > 100) {
          i = 4;
      }
      // Video
      if (m_av_packet->stream_index != *m_videostream) {
        av_packet_free(&m_av_packet);
        m_av_packet = nullptr;
        std::cout << "av_packet free "  << i++ << std::endl;
        return false;
      }
        std::cout << "av_packet not free " << std::endl;
      // Send
      const uint8_t* ptr = m_av_packet->data;
      size_t size = m_av_packet->size;
      while (size > 5) {
        const uint32_t nal_size = htonl(*reinterpret_cast<const uint32_t*>(ptr));
        ptr += 4;
        size -= 4;
        if (nal_size > size)
        {
          std::cout << "Illegal NAL size " << nal_size << std::endl;
          break;
        }
        // Build mpp frame
//        std::cout << "send frame" << std::endl;
        if (!SendFrame(m_mpi, m_ctx, ptr, nal_size, m_packet))
        {
          std::cout << "Failed to send frame: " << nal_size << std::endl;
          exit(1);
        }
        ptr += nal_size;
        size -= nal_size;
      }
      av_packet_free(&m_av_packet);
//      std::cout << "av_packet free 3" << std::endl;
    } else {
//        std::cout << "av package is not null" << std::endl;
    }
    return true;
}

bool Decoder::receive_frame(FrameStruct& frame)
{
    MppFrame source_frame = nullptr;
    // Collect any output frames
    int ret = m_mpi->decode_get_frame(m_ctx, &source_frame);
    if (ret != MPP_OK) {
      std::cout << "Failed to get frame: " << ret << std::endl;
      exit(1);
    }
    if (!source_frame) {
//      std::cout << "get frame is null " << ret << std::endl;
        return false;
    }
    BOOST_SCOPE_EXIT(source_frame)
    {
        mpp_frame_deinit(&source_frame);
    }
    BOOST_SCOPE_EXIT_END
    if (mpp_frame_get_info_change(source_frame)) {
        std::cout << "Frame dimensions and format changed" << std::endl;
        ret = mpp_buffer_group_get_internal(&m_frm_grp, MPP_BUFFER_TYPE_DRM);
        if (ret) {
            std::cout << "Failed to set buffer group" << std::endl;
            exit(1);
        }
        m_mpi->control(m_ctx, MPP_DEC_SET_EXT_BUF_GROUP, m_frm_grp);
        m_mpi->control(m_ctx, MPP_DEC_SET_INFO_CHANGE_READY, nullptr);
        return false;
    }
    frame.mpp_buffer = mpp_frame_get_buffer(source_frame);
    if (frame.mpp_buffer == nullptr) {
        std::cout << "Failed to retrieve buffer from frame" << std::endl;
        exit(1);
    }
    const MppFrameFormat format = mpp_frame_get_fmt(source_frame);
    if (format != MPP_FMT_YUV420SP) {
        std::cout << "Invalid frame format" << std::endl;
        exit(1);
    }

    frame.fd = mpp_buffer_get_fd(frame.mpp_buffer);
    frame.mpp_colour_space = mpp_frame_get_colorspace(source_frame);
    frame.mpp_colour_range = mpp_frame_get_color_range(source_frame);
    frame.mpp_colour_primaries = mpp_frame_get_color_primaries(source_frame);
    frame.width = mpp_frame_get_width(source_frame);
    frame.height = mpp_frame_get_height(source_frame);
    frame.offset_x = mpp_frame_get_offset_x(source_frame);
    frame.offset_y = mpp_frame_get_offset_y(source_frame);
    frame.hor_stride = mpp_frame_get_hor_stride(source_frame);
    frame.ver_stride = mpp_frame_get_ver_stride(source_frame);
    return true;
}

bool Decoder::init_mpp()
{
    int ret = mpp_create(&m_ctx, &m_mpi);
    if (ret != MPP_OK) {
      std::cout << "Failed to create MPP context" << std::endl;
      return false;
    }
    MpiCmd mpi_cmd = MPP_DEC_SET_PARSER_SPLIT_MODE;
    RK_U32 need_split = 1;
    MppParam param = &need_split;
    ret = m_mpi->control(m_ctx, mpi_cmd, param);
    if (ret != MPP_OK) {
      std::cout << "Failed to set MPP split mode" << std::endl;
      return false;
    }
    ret = mpp_init(m_ctx, MPP_CTX_DEC, MPP_VIDEO_CodingAVC);
    if (ret != MPP_OK) {
      std::cout << "Failed to set MPP H264" << std::endl;
      return false;
    }
    return true;
}

bool Decoder::CopyBuffer(const uint8_t *ptr, const size_t size, MppPacket &packet)
{
    const size_t nal_size = size + sizeof(H264_START_SEQUENCE);
    if ((m_packet_buffer == nullptr) || (nal_size > m_packet_buffer_size))
    {
      m_packet_buffer = std::make_unique<char[]>(nal_size);
      m_packet_buffer_size = nal_size;
      if (mpp_packet_deinit(&packet) != MPP_OK)
      {
        std::cout << "Failed to deinit packet" << std::endl;
        return false;
      }
      if (mpp_packet_init(&packet, m_packet_buffer.get(), m_packet_buffer_size) != MPP_OK)
      {
        std::cout << "Failed to init packet" << std::endl;
        return false;
      }
    }
    memcpy(m_packet_buffer.get(), H264_START_SEQUENCE, sizeof(H264_START_SEQUENCE));
    memcpy(m_packet_buffer.get() + 4, ptr, size);
    mpp_packet_write(packet, 0, m_packet_buffer.get(), nal_size);
    mpp_packet_set_pos(packet, m_packet_buffer.get());
    mpp_packet_set_length(packet, nal_size);
    return true;
}

bool Decoder::SendFrame(MppApi *api, MppCtx context, const uint8_t *ptr, const size_t size, MppPacket &packet)
{
    if (!CopyBuffer(ptr, size, packet))
    {
      std::cout << "Failed to copy buffer" << std::endl;
      return false;
    }
    const int ret = api->decode_put_packet(context, packet);
    if (ret != MPP_OK)
    {
      std::cout << "Failed to place packet: " << ret << std::endl;
      return false;
    }
    return true;
}
